from omegaconf import DictConfig

import torch

from ecgcmr.imaging.img_dataset.ImageDataset import BaseImageDataset
from ecgcmr.imaging.img_augmentations.ImageAugmentations import ImageAugmentations


class MaskedImageDataset(BaseImageDataset):
  """
  Dataset that generates data for VideoMAE training
  """
  def __init__(
      self,
      cfg: DictConfig,
      mode: str,
      apply_augmentations: bool = True,
      ) -> None:
    super().__init__(cfg=cfg.dataset.paths, mode=mode)

    self.image_augmentations = ImageAugmentations(cfg=cfg.augmentations.imaging,
                                                  apply_augmentations=apply_augmentations,
                                                  img_size=cfg.dataset.img_size)
    
  def __len__(self):
     return super().__len__()

  def __getitem__(self, index: int) -> torch.Tensor:
    image, es_time_step = super().__getitem__(index)

    image_data = {"image": image, "es_time_step": es_time_step}
    
    return self.image_augmentations(image_data)
